from utils import BloomFilter, CountingBloomFilter
from utils import inet_to_str
import dpkt
import datetime
import socket
from config import Config
from tqdm import tqdm

def get_IP_pkt_info(buf):
    """Get src_IP, src_port, dst_IP, dst_port from buf.

    Args:
        buf: IP raw data.

    Returns:
        src_IP, src_port, dst_IP, dst_port
    """
    # Unpack the IP Packet
    ip_packet = dpkt.ip.IP(buf)

    # Make sure the 'buf' is IP Packet
    if not isinstance(ip_packet, dpkt.ip.IP):
        print('Non IP Packet type not supported %s\n' % buf.__class__.__name__)
        return None
    
    # # Print out IP Packet information
    # print("src IP: ", inet_to_str(ip_packet.src))
    # print("dst IP: ", inet_to_str(ip_packet.dst))

    # Make sure ip_packet.data is TCP Packet
    tcp_packet = ip_packet.data

    if not isinstance(tcp_packet, dpkt.tcp.TCP):
        print('Non TCP Packet type not supported %s\n' % tcp_packet.__class__.__name__)
        return None       

    # # Print out TCP Packet information
    # print("src port: ", tcp_packet.sport)
    # print("dst port: ", tcp_packet.dport)
    return int(ip_packet.src.hex(), 16), tcp_packet.sport, int(ip_packet.dst.hex(), 16), tcp_packet.dport

def get_start_time(pcap_file):
    with open(pcap_file, mode='rb') as f_pcap:
        # Read pcap file
        pcap = dpkt.pcap.Reader(f_pcap)
        return pcap.__next__()[0]

def main():
    LDF = set()

    bf = BloomFilter(error_rate=Config.error_rate, element_num=Config.element_num)
    cbf = CountingBloomFilter(error_rate=Config.error_rate, element_num=Config.element_num)

    start_time = get_start_time(Config.PCAP_FILE_PATH) # start time of a measurement period

    with open(Config.PCAP_FILE_PATH, mode='rb') as f_pcap:
        # Read pcap file
        pcap = dpkt.pcap.Reader(f_pcap)

        for timestamp, buf in tqdm(pcap):
            
            # Case 1: During the a measurement period. 
            if timestamp < start_time + Config.measurement_period_length:

                # Generate ID for each flow;
                flow_id = get_IP_pkt_info(buf)

                if not flow_id:
                    print("Get IP packet information failed!")
                    continue

                # Determine if a flow is LDF.
                bf.add(flow_id)
                if cbf.get_counter_num(flow_id) + 1 >= Config.measurement_times:                    
                    LDF.add(flow_id)

            # Case 2: At the end of a measurement period.
            elif timestamp >= start_time + Config.measurement_period_length:

                # Update start_time;
                start_time += Config.measurement_period_length

                # Update Counting Bloom Filter
                for idx, val in enumerate(bf.bit_array):
                    if val:
                        cbf._set_bit_value(idx, cbf._get_bit_value(idx) + 1)
                    else:
                        cbf._set_bit_value(idx, 0)

                # Update Bloom Filter
                bf.reset()
    return LDF

def output(long_duration_low, output_file_path=Config.result_file_path):
    with open(output_file_path, 'w+') as f_out:
        for flow in long_duration_low:
            # print(list(flow))
            flow = " ".join((str(flow).split(", "))).strip("()")
            f_out.write(flow)
            f_out.write('\n')

def test_output():
    ldf = set([(1928257034, 1116, 985691779, 80), (3260516944, 62447, 2046300540, 23)])
    print(ldf)

if __name__ == '__main__':
    LDF = main()
    print(LDF)

    output(LDF)